"""
债券分析工作流 - LangGraph并行执行架构

使用LangGraph的StateGraph实现完整的债券分析流程：
- 并行执行4个专业分析（信用、条款、宏观、价值）
- 使用get_stream_writer()发送自定义推理事件
- 使用stream_mode="custom"接收实时推理内容
- 已移除冗余的意图识别（由主工作流负责）
"""

import logging
from typing import TypedDict, Annotated, AsyncIterator, Dict, Any, List, Optional
from datetime import datetime
import operator

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.config import get_stream_writer
from langchain_core.runnables import RunnableConfig

from agents.bond_issuer_credit_analysis_agent import BondIssuerCreditAnalysisAgent
from agents.bond_terms_analysis_agent import BondTermsAnalysisAgent
from agents.bond_macro_industry_analysis_agent import BondMacroIndustryAnalysisAgent
from agents.bond_relative_value_analysis_agent import BondRelativeValueAnalysisAgent
from agents.bond_investment_summary_agent import BondInvestmentSummaryAgent


class AnalysisState(TypedDict):
    """债券分析状态"""
    # 输入
    bond_code: str
    user_query: str
    
    # ✅ 新增：分析维度列表（支持部分维度分析）
    analysis_dimensions: List[str]  # ["credit", "terms", "macro", "value"]
    
    # ✅ 分析结果改为单个字符串（每个维度只有一个分析结果，避免累积）
    credit_analysis: str
    terms_analysis: str
    macro_analysis: str
    value_analysis: str
    
    # 最终总结
    final_summary: str
    
    # 元数据
    errors: Annotated[List[str], operator.add]
    completed_steps: Annotated[List[str], operator.add]


class BondAnalysisWorkflow:
    """债券分析工作流"""
    
    def __init__(self):
        """初始化工作流"""
        self.logger = logging.getLogger("workflows.bond_analysis")
        
        # 初始化专业分析智能体
        self.credit_agent = BondIssuerCreditAnalysisAgent()
        self.terms_agent = BondTermsAnalysisAgent()
        self.macro_agent = BondMacroIndustryAnalysisAgent()
        self.value_agent = BondRelativeValueAnalysisAgent()
        self.summary_agent = BondInvestmentSummaryAgent()
        
        # ✅ 移除检查点保存器，避免状态累积
        # 每次分析都是独立的，不会复用之前的状态
        self.checkpointer = None
        
        # 构建工作流图
        self.app = self._build_workflow()
    
    def _build_workflow(self):
        """构建工作流状态图（支持维度选择的条件路由）"""
        # 创建状态图
        workflow = StateGraph(AnalysisState)
        
        # 添加所有节点
        workflow.add_node("credit_analysis", self._credit_analysis_node)
        workflow.add_node("terms_analysis", self._terms_analysis_node)
        workflow.add_node("macro_analysis", self._macro_analysis_node)
        workflow.add_node("value_analysis", self._value_analysis_node)
        workflow.add_node("summary", self._summary_node)
        
        # ✅ 添加dispatch节点（支持维度选择）
        workflow.add_node("dispatch", self._dispatch_node)
        
        # 设置入口点
        workflow.set_entry_point("dispatch")
        
        # ✅ 从dispatch节点使用条件边路由到选中的分析节点
        workflow.add_conditional_edges(
            "dispatch",
            self._route_from_dispatch,
            ["credit_analysis", "terms_analysis", "macro_analysis", "value_analysis"]
        )
        
        # 所有分析完成后汇聚到总结节点
        workflow.add_edge("credit_analysis", "summary")
        workflow.add_edge("terms_analysis", "summary")
        workflow.add_edge("macro_analysis", "summary")
        workflow.add_edge("value_analysis", "summary")
        
        # 总结后结束
        workflow.add_edge("summary", END)
        
        # 编译工作流（不使用checkpointer，确保每次分析独立）
        return workflow.compile()
    
    def _dispatch_node(self, state: AnalysisState) -> Dict[str, Any]:
        """
        分发节点：根据analysis_dimensions决定要执行的分析节点
        
        维度映射：
        - "credit": credit_analysis (发行人信用分析)
        - "terms": terms_analysis (债券条款分析)
        - "macro": macro_analysis (宏观行业分析)
        - "value": value_analysis (相对价值分析)
        """
        dimensions = state.get("analysis_dimensions", [])
        
        # 默认全量分析（如果未指定或维度>=3个）
        if not dimensions or len(dimensions) >= 3:
            dimensions = ["credit", "terms", "macro", "value"]
        
        # 发送推理事件
        writer = get_stream_writer()
        writer({
            "type": "reasoning",
            "node": "dispatch",
            "content": f"📋 分析维度: {', '.join(dimensions)} ({len(dimensions)}个)",
            "timestamp": datetime.now().isoformat()
        })
        
        return {"analysis_dimensions": dimensions}
    
    def _route_from_dispatch(self, state: AnalysisState) -> List[str]:
        """
        路由函数：返回要执行的节点列表
        
        Returns:
            节点名称列表，LangGraph会并行执行这些节点
        """
        dimensions = state.get("analysis_dimensions", [])
        
        # 维度到节点的映射
        dimension_map = {
            "credit": "credit_analysis",
            "terms": "terms_analysis",
            "macro": "macro_analysis",
            "value": "value_analysis"
        }
        
        # 根据维度返回要执行的节点
        nodes = [dimension_map[dim] for dim in dimensions if dim in dimension_map]
        
        # 至少执行一个节点（默认信用分析）
        if not nodes:
            nodes = ["credit_analysis"]
        
        self.logger.info(f"路由到节点: {nodes}")
        return nodes
    
    async def _credit_analysis_node(self, state: AnalysisState) -> Dict[str, Any]:
        """发行人信用分析节点 - 使用自定义事件流"""
        try:
            self.logger.info(f"开始信用分析: {state['bond_code']}")
            
            # ✅ 获取stream writer
            writer = get_stream_writer()
            
            # 发送节点开始事件
            writer({
                "type": "step_start",
                "node": "credit_analysis",
                "step_name": "🏦 发行人信用分析",
                "timestamp": datetime.now().isoformat()
            })
            
            results = []
            
            # ✅ 流式执行智能体，转发所有推理事件
            async for event in self.credit_agent.stream_analyze(state['bond_code']):
                event_type = event.get("type")
                
                if event_type == "reasoning":
                    # ✅ 转发所有推理内容，不再过滤
                    content = event.get("content", "")
                    if content:
                        writer({
                            "type": "reasoning",
                            "node": "credit_analysis",
                            "content": content,
                            "timestamp": datetime.now().isoformat()
                        })
                
                elif event_type == "final_answer":
                    # 只保存结果，不发送到UI（避免在Step中显示完整结果）
                    results.append(event.get("content", ""))
            
            # 发送节点结束事件
            writer({
                "type": "step_end",
                "node": "credit_analysis",
                "step_name": "🏦 发行人信用分析",
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "credit_analysis": "\n".join(results),  # ✅ 拼接为单个字符串
                "completed_steps": ["credit_analysis"]
            }
        except Exception as e:
            self.logger.error(f"信用分析失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "credit_analysis",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "credit_analysis": [f"信用分析失败: {str(e)}"],
                "errors": [f"信用分析失败: {str(e)}"],
                "completed_steps": ["credit_analysis"]
            }
    
    async def _terms_analysis_node(self, state: AnalysisState) -> Dict[str, Any]:
        """债券条款分析节点 - 使用自定义事件流"""
        try:
            self.logger.info(f"开始条款分析: {state['bond_code']}")
            
            # ✅ 获取stream writer
            writer = get_stream_writer()
            
            # 发送节点开始事件
            writer({
                "type": "step_start",
                "node": "terms_analysis",
                "step_name": "📜 债券条款分析",
                "timestamp": datetime.now().isoformat()
            })
            
            results = []
            
            # ✅ 流式执行智能体，转发所有推理事件
            async for event in self.terms_agent.stream_analyze(state['bond_code']):
                event_type = event.get("type")
                
                if event_type == "reasoning":
                    # ✅ 转发所有推理内容，不再过滤
                    content = event.get("content", "")
                    if content:
                        writer({
                            "type": "reasoning",
                            "node": "terms_analysis",
                            "content": content,
                            "timestamp": datetime.now().isoformat()
                        })
                
                elif event_type == "final_answer":
                    # 只保存结果，不发送到UI
                    results.append(event.get("content", ""))
            
            # 发送节点结束事件
            writer({
                "type": "step_end",
                "node": "terms_analysis",
                "step_name": "📜 债券条款分析",
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "terms_analysis": "\n".join(results),  # ✅ 拼接为单个字符串
                "completed_steps": ["terms_analysis"]
            }
        except Exception as e:
            self.logger.error(f"条款分析失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "terms_analysis",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "terms_analysis": [f"条款分析失败: {str(e)}"],
                "errors": [f"条款分析失败: {str(e)}"],
                "completed_steps": ["terms_analysis"]
            }
    
    async def _macro_analysis_node(self, state: AnalysisState) -> Dict[str, Any]:
        """宏观行业分析节点 - 使用自定义事件流"""
        try:
            self.logger.info(f"开始宏观分析: {state['bond_code']}")
            
            # ✅ 获取stream writer
            writer = get_stream_writer()
            
            # 发送节点开始事件
            writer({
                "type": "step_start",
                "node": "macro_analysis",
                "step_name": "🌍 宏观行业分析",
                "timestamp": datetime.now().isoformat()
            })
            
            results = []
            
            # ✅ 流式执行智能体，转发所有推理事件
            async for event in self.macro_agent.stream_analyze(state['bond_code']):
                event_type = event.get("type")
                
                if event_type == "reasoning":
                    # ✅ 转发所有推理内容，不再过滤
                    content = event.get("content", "")
                    if content:
                        writer({
                            "type": "reasoning",
                            "node": "macro_analysis",
                            "content": content,
                            "timestamp": datetime.now().isoformat()
                        })
                
                elif event_type == "final_answer":
                    # 只保存结果，不发送到UI
                    results.append(event.get("content", ""))
            
            # 发送节点结束事件
            writer({
                "type": "step_end",
                "node": "macro_analysis",
                "step_name": "🌍 宏观行业分析",
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "macro_analysis": "\n".join(results),  # ✅ 拼接为单个字符串
                "completed_steps": ["macro_analysis"]
            }
        except Exception as e:
            self.logger.error(f"宏观分析失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "macro_analysis",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "macro_analysis": [f"宏观分析失败: {str(e)}"],
                "errors": [f"宏观分析失败: {str(e)}"],
                "completed_steps": ["macro_analysis"]
            }
    
    async def _value_analysis_node(self, state: AnalysisState) -> Dict[str, Any]:
        """相对价值分析节点 - 使用自定义事件流"""
        try:
            self.logger.info(f"开始价值分析: {state['bond_code']}")
            
            # ✅ 获取stream writer
            writer = get_stream_writer()
            
            # 发送节点开始事件
            writer({
                "type": "step_start",
                "node": "value_analysis",
                "step_name": "💎 相对价值分析",
                "timestamp": datetime.now().isoformat()
            })
            
            results = []
            
            # ✅ 流式执行智能体，转发所有推理事件
            async for event in self.value_agent.stream_analyze(state['bond_code']):
                event_type = event.get("type")
                
                if event_type == "reasoning":
                    # ✅ 转发所有推理内容，不再过滤
                    content = event.get("content", "")
                    if content:
                        writer({
                            "type": "reasoning",
                            "node": "value_analysis",
                            "content": content,
                            "timestamp": datetime.now().isoformat()
                        })
                
                elif event_type == "final_answer":
                    # 只保存结果，不发送到UI
                    results.append(event.get("content", ""))
            
            # 发送节点结束事件
            writer({
                "type": "step_end",
                "node": "value_analysis",
                "step_name": "💎 相对价值分析",
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "value_analysis": "\n".join(results),  # ✅ 拼接为单个字符串
                "completed_steps": ["value_analysis"]
            }
        except Exception as e:
            self.logger.error(f"价值分析失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "value_analysis",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "value_analysis": [f"价值分析失败: {str(e)}"],
                "errors": [f"价值分析失败: {str(e)}"],
                "completed_steps": ["value_analysis"]
            }
    
    async def _summary_node(self, state: AnalysisState) -> Dict[str, Any]:
        """投资总结节点 - 使用自定义事件流"""
        try:
            self.logger.info(f"开始投资总结: {state['bond_code']}")
            
            # ✅ 获取stream writer
            writer = get_stream_writer()
            
            # 发送节点开始事件
            writer({
                "type": "step_start",
                "node": "summary",
                "step_name": "📊 投资总结",
                "timestamp": datetime.now().isoformat()
            })
            
            # 等待所有分析完成
            required_steps = ["credit_analysis", "terms_analysis", "macro_analysis", "value_analysis"]
            completed = state.get("completed_steps", [])
            
            if not all(step in completed for step in required_steps):
                self.logger.warning(f"部分分析未完成，已完成: {completed}")
            
            # ✅ 收集其他维度的分析结果（现在直接是字符串）
            analysis_context = {
                "bond_code": state['bond_code'],
                "credit_analysis": state.get("credit_analysis", ""),
                "terms_analysis": state.get("terms_analysis", ""),
                "macro_analysis": state.get("macro_analysis", ""),
                "value_analysis": state.get("value_analysis", "")
            }
            
            results = []
            
            # ✅ 流式执行智能体，传递其他维度的分析结果作为上下文
            async for event in self.summary_agent.stream_analyze_with_context(analysis_context):
                event_type = event.get("type")
                
                if event_type == "reasoning":
                    # ✅ 转发所有推理内容，不再过滤
                    content = event.get("content", "")
                    if content:
                        writer({
                            "type": "reasoning",
                            "node": "summary",
                            "content": content,
                            "timestamp": datetime.now().isoformat()
                        })
                
                elif event_type == "final_answer":
                    # 只保存结果，不发送到UI
                    results.append(event.get("content", ""))
            
            summary = "\n".join(results) if results else "投资总结生成失败"
            
            # 发送节点结束事件
            writer({
                "type": "step_end",
                "node": "summary",
                "step_name": "📊 投资总结",
                "timestamp": datetime.now().isoformat()
            })
            
            # ✅ 发送final_result事件，包含完整的分析状态（包括summary）
            # 显式构建完整状态，确保所有字段都包含
            complete_state = {
                "bond_code": state.get("bond_code", ""),
                "user_query": state.get("user_query", ""),
                "analysis_dimensions": state.get("analysis_dimensions", []),
                "credit_analysis": state.get("credit_analysis", ""),  # ✅ 字符串
                "terms_analysis": state.get("terms_analysis", ""),    # ✅ 字符串
                "macro_analysis": state.get("macro_analysis", ""),    # ✅ 字符串
                "value_analysis": state.get("value_analysis", ""),    # ✅ 字符串
                "final_summary": summary,  # ← 确保包含summary
                "errors": state.get("errors", []),
                "completed_steps": state.get("completed_steps", []) + ["summary"]
            }
            
            self.logger.info(f"发送final_result事件，final_summary长度: {len(summary)}")
            
            writer({
                "type": "final_result",
                "state": complete_state,  # 传递包含所有字段的完整状态
                "timestamp": datetime.now().isoformat()
            })
            
            return {
                "final_summary": summary,
                "completed_steps": ["summary"]
            }
        except Exception as e:
            self.logger.error(f"投资总结失败: {e}")
            writer = get_stream_writer()
            writer({
                "type": "error",
                "node": "summary",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            })
            return {
                "final_summary": f"投资总结失败: {str(e)}",
                "errors": [f"投资总结失败: {str(e)}"],
                "completed_steps": ["summary"]
            }
    
    async def execute(
        self, 
        bond_code: str, 
        user_query: str = "",
        thread_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        执行债券分析工作流
        
        Args:
            bond_code: 债券代码
            user_query: 用户查询（可选）
            thread_id: 线程ID（用于状态持久化）
            
        Returns:
            完整的分析结果
        """
        if not thread_id:
            thread_id = f"bond_analysis_{bond_code}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        # 初始化状态
        initial_state: AnalysisState = {
            "bond_code": bond_code,
            "user_query": user_query or f"分析债券{bond_code}",
            "analysis_dimensions": [],  # ✅ 添加维度字段（空列表=全量分析）
            "credit_analysis": "",  # ✅ 字符串
            "terms_analysis": "",   # ✅ 字符串
            "macro_analysis": "",   # ✅ 字符串
            "value_analysis": "",   # ✅ 字符串
            "final_summary": "",
            "errors": [],
            "completed_steps": []
        }
        
        try:
            # 执行工作流
            config = RunnableConfig(configurable={"thread_id": thread_id})
            final_state = await self.app.ainvoke(initial_state, config)
            
            # 格式化结果
            return {
                "status": "success",
                "bond_code": bond_code,
                "timestamp": datetime.now().isoformat(),
                "results": {
                    "credit_analysis": final_state.get("credit_analysis", ""),  # ✅ 直接使用字符串
                    "terms_analysis": final_state.get("terms_analysis", ""),    # ✅ 直接使用字符串
                    "macro_analysis": final_state.get("macro_analysis", ""),    # ✅ 直接使用字符串
                    "value_analysis": final_state.get("value_analysis", ""),    # ✅ 直接使用字符串
                    "final_summary": final_state.get("final_summary", "")
                },
                "metadata": {
                    "completed_steps": final_state.get("completed_steps", []),
                    "errors": final_state.get("errors", [])
                }
            }
        except Exception as e:
            self.logger.error(f"工作流执行失败: {e}")
            return {
                "status": "error",
                "bond_code": bond_code,
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
    
    async def stream_execute(
        self,
        bond_code: str,
        user_query: str = "",
        thread_id: Optional[str] = None,
        analysis_dimensions: Optional[List[str]] = None
    ) -> AsyncIterator[Dict[str, Any]]:
        """
        流式执行债券分析工作流（使用LangGraph自定义事件流）
        
        ✅ 方案3实现：
        - 使用StateGraph定义的并行边（真正的LangGraph并行）
        - 使用get_stream_writer()在节点内发送自定义事件
        - 使用stream_mode="custom"接收实时推理内容
        
        Args:
            bond_code: 债券代码
            user_query: 用户查询
            thread_id: 线程ID
            
        Yields:
            工作流执行的每个步骤事件，包括智能体的实时推理内容
        """
        if not thread_id:
            thread_id = f"bond_analysis_{bond_code}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        # 初始化状态
        initial_state: AnalysisState = {
            "bond_code": bond_code,
            "user_query": user_query or f"分析债券{bond_code}",
            "analysis_dimensions": analysis_dimensions or [],  # ✅ 支持指定维度
            "credit_analysis": "",  # ✅ 字符串
            "terms_analysis": "",   # ✅ 字符串
            "macro_analysis": "",   # ✅ 字符串
            "value_analysis": "",   # ✅ 字符串
            "final_summary": "",
            "errors": [],
            "completed_steps": []
        }
        
        config = RunnableConfig(configurable={"thread_id": thread_id})
        
        try:
            # ✅ 使用stream_mode="custom"接收节点内通过writer发送的自定义事件
            final_state = None
            async for event in self.app.astream(
                initial_state,
                config,
                stream_mode="custom"  # 关键：只接收自定义事件
            ):
                # 直接转发自定义事件给UI层
                yield event
                
                # ✅ 捕获final_result事件中的状态
                if event.get("type") == "final_result":
                    final_state = event.get("state", {})
            
            # ✅ 如果没有从事件中获取到状态，发送完成事件
            if not final_state:
                # 由于没有checkpointer，我们通过initial_state构建一个基本的最终状态
                # 注意：实际结果已通过final_result事件发送
                yield {
                    "type": "workflow_complete",
                    "state": initial_state,  # 基本状态
                    "timestamp": datetime.now().isoformat()
                }
            
        except Exception as e:
            self.logger.error(f"流式执行失败: {e}", exc_info=True)
            yield {
                "type": "error",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }


# 全局实例
bond_analysis_workflow = BondAnalysisWorkflow()


# 便捷函数
async def analyze_bond(bond_code: str, user_query: str = "") -> Dict[str, Any]:
    """
    分析债券的便捷函数
    
    Args:
        bond_code: 债券代码
        user_query: 用户查询
        
    Returns:
        分析结果
    """
    return await bond_analysis_workflow.execute(bond_code, user_query)